# semantic_comparer.py
# -*- coding: utf-8 -*-

import logging
import time
from typing import List, Tuple, Dict, Any, Optional, TYPE_CHECKING
import torch
import config
import plagiarism_core.preprocessor as preprocessor # 导入 preprocessor 用于语料库文本分块

# 配置日志记录器
logger = logging.getLogger(__name__)

# 导入 sentence-transformers 相关库
SENTENCE_TRANSFORMERS_AVAILABLE = False
util = None
try:
    from sentence_transformers import SentenceTransformer, util
    SENTENCE_TRANSFORMERS_AVAILABLE = True
    logger.info("成功导入 sentence-transformers 库。")
except ImportError:
    logger.warning("semantic_comparer: 未找到 'sentence-transformers' 库。语义相似度功能将不可用。")
    # 定义一个模拟的 util 对象，以便在库不可用时代码也能运行（但功能受限）
    class MockUtil:
        def semantic_search(*args, **kwargs):
            logger.warning("模拟 semantic_search 被调用，返回空列表。")
            return []
        def cos_sim(*args, **kwargs):
            logger.warning("模拟 cos_sim 被调用，返回零张量。")
            return torch.tensor([[0.0]]) # 返回一个符合类型的张量
    util = MockUtil()

# 仅用于类型检查时导入，帮助静态分析器识别类型
if TYPE_CHECKING:
    from sentence_transformers import SentenceTransformer

# 定义语义匹配结果的数据结构
# 元组包含:
# - input_chunk_index (int): 输入文本中块的索引。
# - source_corpus_id (str): 来源语料库文档的文件名或标识符。
# - source_chunk_index (int): 来源语料库文档中匹配块的索引 (在该文档内的索引)。
# - similarity_score (float): 计算得到的语义相似度分数（通常是余弦相似度）。
# - source_start_char (int): 匹配块在来源文档 *清理后* 文本中的起始字符索引。
# - source_end_char (int): 匹配块在来源文档 *清理后* 文本中的结束字符索引（不包含）。
SemanticMatchWithPosTuple = Tuple[int, str, int, float, int, int]

class SemanticComparer:
    """
    封装语义相似度计算逻辑的工具类。

    负责加载语义模型，并提供计算文本块之间相似度的方法。
    """

    def __init__(self):
        """
        初始化语义比较器。

        根据配置加载 Sentence Transformer 模型，并自动检测可用设备。
        """
        # 从配置中读取相关参数
        self.chunk_type: str = config.SEMANTIC_CHUNK_TYPE
        self.threshold: float = config.SEMANTIC_SIMILARITY_THRESHOLD
        self.model: Optional[SentenceTransformer] = None # 类型明确为 SentenceTransformer 或 None
        self.enabled: bool = False # 标记功能是否启用
        self.device: Optional[str] = None  # 存储检测到的设备 ('cuda', 'mps', 'cpu')

        # 检查配置是否启用语义相似度
        if config.USE_SEMANTIC_SIMILARITY:
            if SENTENCE_TRANSFORMERS_AVAILABLE:
                # --- 自动检测最佳可用设备 ---
                if torch.cuda.is_available():
                    self.device = 'cuda'
                elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): # 检查 MPS (Apple Silicon GPU)
                    # 注意：MPS 的支持可能取决于 PyTorch 和 sentence-transformers 的版本
                    try:
                        # 尝试创建一个 MPS 张量来确认可用性
                        _ = torch.tensor([1], device='mps')
                        self.device = 'mps'
                    except Exception:
                         logger.warning("检测到 MPS 后端但无法创建 MPS 张量，将回退到 CPU。")
                         self.device = 'cpu'
                else:
                    self.device = 'cpu'
                logger.info(f"SemanticComparer: 自动检测到可用设备: {self.device}")
                # -----------------------------

                # --- 加载模型 ---
                model_name = config.SENTENCE_MODEL_NAME
                try:
                    logger.info(f"SemanticComparer: 开始加载语义模型 '{model_name}' 到设备 '{self.device}'...")
                    start_time = time.time()
                    # 加载模型，明确指定设备
                    # trust_remote_code=True 可能对某些模型是必需的，但有安全风险，谨慎使用
                    self.model = SentenceTransformer(model_name, device=self.device, trust_remote_code=True)
                    end_time = time.time()
                    logger.info(f"SemanticComparer: 语义模型 '{model_name}' 加载成功，耗时: {end_time - start_time:.2f} 秒。")
                    self.enabled = True # 标记为启用
                except Exception as e:
                    logger.exception(f"SemanticComparer: 加载语义模型 '{model_name}' 失败: {e}")
                    self.model = None
                    self.enabled = False
                    self.device = None # 加载失败则设备也重置
            else:
                logger.warning("SemanticComparer: 'sentence-transformers' 库未安装，无法启用语义相似度功能。")
                self.enabled = False
        else:
            logger.info("SemanticComparer: 配置中已禁用语义相似度计算 (USE_SEMANTIC_SIMILARITY = False)。")
            self.enabled = False

    def find_similar_chunks(self,
                            input_chunks_with_pos: List[Tuple[str, int, int]],
                            corpus_texts_to_compare: Dict[str, str]
                           ) -> List[SemanticMatchWithPosTuple]:
        """
        计算给定输入文本块与指定语料库文本之间的语义相似块。

        此方法会对 `corpus_texts_to_compare` 中的文本进行实时分块和嵌入，
        然后使用 `semantic_search` 查找与输入块相似的语料块。
        返回的匹配结果包含了来源块在其原始文档（清理后）中的字符位置信息。
        """

        logger.info("\n")
        logger.info("------- 开始计算语义相似度（详细比较阶段）-------")
        # 检查功能是否启用和模型是否加载
        if not self.enabled or not self.model:
            logger.warning("语义相似度计算已禁用或模型未加载，无法执行 find_similar_chunks。")
            return []
        # 检查输入是否有效
        if not input_chunks_with_pos:
            logger.warning("输入的文本块列表为空，无法进行语义比较。")
            return []
        if not corpus_texts_to_compare:
            logger.warning("没有提供需要比较的语料库文本 (corpus_texts_to_compare 为空)，跳过语义比较。")
            return []

        logger.info(f"开始计算语义相似度。输入块数量: {len(input_chunks_with_pos)}, 待比较语料库文档数量: {len(corpus_texts_to_compare)}")
        logger.info(f"使用分块方式: '{self.chunk_type}', 相似度阈值: {self.threshold:.2f}")
        start_time = time.time()

        # --- 1. 准备输入块文本 ---
        input_chunks_text = [item[0] for item in input_chunks_with_pos]

        # --- 2. 对需要比较的语料库文本进行动态分块 ---
        #    这一步对于没有使用语义索引，或者需要对索引结果进行精确重排/验证时是必要的。
        corpus_chunks_info: List[Dict[str, Any]] = []  # 存储所有语料块的详细信息
        global_corpus_chunk_index = 0 # 用于 semantic_search 返回的 corpus_id (实际是索引)
        logger.info(f"正在对 {len(corpus_texts_to_compare)} 个候选语料库文档进行动态分块 (方式: {self.chunk_type})...")
        corpus_chunking_start_time = time.time()

        for corpus_id, corpus_text in corpus_texts_to_compare.items():
            if not corpus_text or not corpus_text.strip():
                logger.debug(f"跳过空的语料库文档进行分块: {corpus_id}")
                continue

            # --- 调用 preprocessor 获取带位置的块 ---
            # 这里需要确保 preprocessor 可用且能正常工作
            corpus_chunks_for_this_doc: List[Tuple[str, int, int]] = []
            try:
                # 根据配置选择分块函数
                if self.chunk_type == 'sentence':
                    # 调用分句函数，它内部会处理长度过滤和长句切分
                    # 注意：这里假设 preprocessor 已加载 LTP 或有回退机制
                    corpus_chunks_for_this_doc = preprocessor.split_into_sentences_with_pos(corpus_text)
                elif self.chunk_type == 'paragraph':
                    corpus_chunks_for_this_doc = preprocessor.split_into_paragraphs_with_pos(corpus_text)
                else: # 处理未知的 chunk_type
                    logger.warning(f"配置了未知的 SEMANTIC_CHUNK_TYPE '{self.chunk_type}'，将默认使用句子分块。")
                    corpus_chunks_for_this_doc = preprocessor.split_into_sentences_with_pos(corpus_text)

                # --- 存储块信息 ---
                if corpus_chunks_for_this_doc:
                    logger.debug(f"文档 '{corpus_id}' 分割出 {len(corpus_chunks_for_this_doc)} 个有效块。")
                    for i, (chunk_text, start_char, end_char) in enumerate(corpus_chunks_for_this_doc):
                        # chunk_text 应该是 preprocessor 返回时已经 strip 过的
                        chunk_info = {
                            "global_index": global_corpus_chunk_index, # 全局索引，用于 semantic_search
                            "text": chunk_text,              # 块文本
                            "corpus_id": corpus_id,          # 所属原始文档 ID
                            "corpus_chunk_index": i,         # 在本文档内的块索引
                            "start_char": start_char,        # 在本文档清理后文本中的起始位置
                            "end_char": end_char             # 在本文档清理后文本中的结束位置
                        }
                        corpus_chunks_info.append(chunk_info)
                        global_corpus_chunk_index += 1
                else:
                     logger.debug(f"文档 '{corpus_id}' 未分割出有效的块。")

            except Exception as e:
                 # 处理分块过程中可能发生的错误
                 logger.error(f"处理语料库文档 '{corpus_id}' 进行分块时发生错误: {e}", exc_info=True)
                 # 可以选择跳过这个文档或采取其他错误处理

        chunking_duration = time.time() - corpus_chunking_start_time
        logger.info(f"候选语料库分块完成，耗时 {chunking_duration:.2f} 秒。共生成 {len(corpus_chunks_info)} 个语料块。")
        # 如果没有有效的语料块，无法进行比较
        if not corpus_chunks_info:
            logger.warning("候选语料库未产生任何有效分块，无法进行语义比较。")
            return []
        # -------------------------------------------------------------

        # 提取所有语料块的文本用于嵌入
        corpus_chunks_text = [item["text"] for item in corpus_chunks_info]

        # --- 3. 计算向量 (Embeddings) ---
        embedding_start_time = time.time()
        try:
            logger.info(f"开始计算输入块 ({len(input_chunks_text)}) 的向量...")
            input_embeddings = self.model.encode(
                input_chunks_text,
                convert_to_tensor=True,
                show_progress_bar=False, # 在日志中显示进度可能更好
                device=self.device,      # 明确指定设备
                batch_size=config.SEMANTIC_ENCODE_BATCH_SIZE # 使用配置的批量大小
            )
            logger.info(f"输入块向量计算完成。")

            logger.info(f"开始计算语料库块 ({len(corpus_chunks_text)}) 的向量...")
            corpus_embeddings = self.model.encode(
                corpus_chunks_text,
                convert_to_tensor=True,
                show_progress_bar=False,
                device=self.device,
                batch_size=config.SEMANTIC_ENCODE_BATCH_SIZE
            )
            logger.info(f"语料库块向量计算完成。")
        except Exception as e:
            logger.exception(f"计算文本向量 (Embeddings) 时发生严重错误: {e}")
            return [] # 嵌入失败则无法继续

        embedding_duration = time.time() - embedding_start_time
        logger.info(f"向量计算完成，耗时 {embedding_duration:.2f} 秒。")

        # --- 4. 计算相似度并筛选 ---
        semantic_matches_with_pos: List[SemanticMatchWithPosTuple] = []
        search_start_time = time.time()
        try:
            logger.info("开始执行语义搜索 (查找相似块)...")
            # 使用 semantic_search 函数查找最相似的 top_k 个语料块
            # top_k 的值可以从配置中获取，或者设置为一个合理的默认值
            # 注意：如果语料块数量小于 top_k，它会返回所有语料块的得分
            top_k_search = config.SEMANTIC_SEARCH_TOP_K
            # 确保 top_k 不大于实际的语料块数量，且至少为 1
            top_k_search = min(top_k_search, len(corpus_chunks_info)) if len(corpus_chunks_info) > 0 else 1
            if top_k_search <= 0: top_k_search = 1
            logger.debug(f"Semantic search top_k 设置为: {top_k_search}")

            # score_function=util.cos_sim 指定使用余弦相似度
            # 返回结果 hits 是一个列表，长度等于输入块的数量
            # 每个元素是一个列表，包含 top_k 个最相似的语料块信息 (字典)
            hits = util.semantic_search(
                input_embeddings,
                corpus_embeddings,
                top_k=top_k_search,
                score_function=util.cos_sim # 明确使用余弦相似度
            )
            search_duration = time.time() - search_start_time
            logger.info(f"语义搜索完成，耗时 {search_duration:.2f} 秒。开始筛选结果...")

            processed_hits_count = 0
            final_match_count = 0
            # 遍历每个输入块的搜索结果
            for input_chunk_idx, chunk_hits in enumerate(hits):
                # 获取当前输入块的信息（用于日志或调试）
                input_chunk_text, input_start, input_end = input_chunks_with_pos[input_chunk_idx]

                # 遍历当前输入块找到的 top_k 个最相似语料块
                for hit in chunk_hits:
                    corpus_global_idx = hit['corpus_id'] # 注意：这里的 'corpus_id' 是 corpus_chunks_info 列表的索引
                    score = hit['score']
                    processed_hits_count += 1

                    # 应用相似度阈值进行过滤
                    # 由于 semantic_search 返回的结果是按得分降序排列的，
                    # 一旦遇到低于阈值的，后续的得分只会更低，可以提前终止内层循环
                    if score < self.threshold:
                        break # 跳出当前输入块的内层循环

                    # --- 验证索引并获取来源块的详细信息 ---
                    if 0 <= corpus_global_idx < len(corpus_chunks_info):
                        # 从我们之前构建的 corpus_chunks_info 列表中获取详细信息
                        corpus_chunk_detail = corpus_chunks_info[corpus_global_idx]
                        original_corpus_id = corpus_chunk_detail["corpus_id"] # 真实的语料库文档 ID
                        corpus_chunk_idx_in_doc = corpus_chunk_detail["corpus_chunk_index"] # 在其源文档中的块索引
                        src_start_char = corpus_chunk_detail["start_char"]
                        src_end_char = corpus_chunk_detail["end_char"]
                        src_chunk_text = corpus_chunk_detail["text"] # 用于日志

                        logger.debug(f"  [语义匹配] 输入块 {input_chunk_idx} ({input_start}-{input_end}) "
                                     f"与 语料库文档 '{original_corpus_id}' -> 块 {corpus_chunk_idx_in_doc} ({src_start_char}-{src_end_char}) "
                                     f"得分: {score:.4f} (阈值: {self.threshold:.2f})")
                        # 可以取消下面的注释以打印具体文本片段进行调试
                        # logger.debug(f"    输入块内容: '{repr(input_chunk_text[:80])}...'")
                        # logger.debug(f"    来源块内容: '{repr(src_chunk_text[:80])}...'")

                        # 构建包含来源位置信息的 SemanticMatchWithPosTuple
                        match_info: SemanticMatchWithPosTuple = (
                            input_chunk_idx,         # 输入块索引
                            original_corpus_id,      # 来源语料库文档 ID
                            corpus_chunk_idx_in_doc, # 来源块在其文档内的索引
                            score,                   # 相似度得分
                            src_start_char,          # 来源块起始位置
                            src_end_char             # 来源块结束位置
                        )
                        semantic_matches_with_pos.append(match_info)
                        final_match_count += 1
                    else:
                        # 如果 semantic_search 返回的索引越界，记录错误
                        logger.error(f"语义搜索返回了无效的语料库块索引: {corpus_global_idx} (有效范围 0-{len(corpus_chunks_info)-1})")
                        # 可以考虑在这里添加更健壮的处理，但通常不应发生

            logger.info(f"语义搜索结果筛选完成。共检查了 {processed_hits_count} 个候选匹配项。")

        except Exception as e:
            logger.exception(f"执行语义搜索或筛选结果时发生错误: {e}")
            return []  # 出错则返回空列表

        total_duration = time.time() - start_time
        logger.info(f"------- 语义相似度计算（详细比较阶段）完成，共找到 {final_match_count} 个高于阈值的匹配，总耗时: {total_duration:.2f} 秒。-------")
        return semantic_matches_with_pos


    def test_similarity(self, text1: str, text2: str) -> Optional[float]:
        """
        单步测试方法：计算两个给定文本片段的语义相似度得分。
        用于快速验证模型对特定文本对的相似度判断。
        """
        if not self.enabled or not self.model:
            print("【测试】SemanticComparer 未启用或模型未加载，无法进行相似度测试。")
            logger.warning("test_similarity: SemanticComparer 未启用或模型未加载。")
            return None

        if not text1 or not text2:
            print("【测试】错误：输入的文本片段不能为空。")
            logger.warning("test_similarity: 输入文本为空。")
            return None

        print("\n--- 语义相似度单步测试 ---")
        print(f"文本 1: {repr(text1)}")
        print(f"文本 2: {repr(text2)}")
        logger.debug(f"开始测试相似度: text1='{repr(text1)}', text2='{repr(text2)}'")

        try:
            start_time = time.time()
            # 1. 将两个文本编码为向量
            #    注意：即使是单个文本，encode 也期望输入是列表
            embeddings = self.model.encode(
                [text1, text2],
                convert_to_tensor=True,
                show_progress_bar=False, # 测试时通常不需要进度条
                device=self.device
            )

            # 2. 计算余弦相似度
            cosine_scores = util.cos_sim(embeddings[0].unsqueeze(0), embeddings[1].unsqueeze(0))

            # 从返回的 [[score]] 张量中提取分数
            # .item() 将只有一个元素的张量转换为 Python float 类型
            similarity_score = cosine_scores[0][0].item()

            end_time = time.time()

            print(f"计算出的相似度得分: {similarity_score:.4f}")
            print(f"计算耗时: {end_time - start_time:.4f} 秒")
            logger.debug(f"测试相似度结果: {similarity_score:.4f}, 耗时: {end_time - start_time:.4f} 秒")
            print("--- 测试结束 ---")
            return similarity_score

        except Exception as e:
            print(f"【测试】计算相似度时发生错误: {e}")
            logger.exception("test_similarity 执行过程中发生错误:")
            print("--- 测试结束 ---")
            return None
